package com.lx.algorithm.Tree;


/**
 * Description:
 * Copyright:   Copyright (c)2019
 * Company:     zefu
 *
 * @author: 张李鑫
 * @version: 1.0
 * Create at:   2021-11-01 13:49:06
 * <p>
 * Modification History:
 * Date         Author      Version     Description
 * ------------------------------------------------------------------
 * 2021-11-01     张李鑫                     1.0         1.0 Version
 */
public class AVLTree {

    public static class AVLNode<K extends Comparable<K>, V> {
        public K k;
        public V v;
        public AVLNode<K, V> l;
        public AVLNode<K, V> r;
        public int h;

        public AVLNode(K key, V value) {
            k = key;
            v = value;
            h = 1;
        }
    }

    public static class AVLTreeMap<K extends Comparable<K>, V> {
        private AVLNode<K, V> root;
        private int size;

        public AVLTreeMap() {
            root = null;
            size = 0;
        }

        /**
         * <p>
         * 右旋
         * </p>
         *
         * @param cur
         * @return com.lx.algorithm.Tree.AVLTree.AVLNode<K, V>
         * @Author zlx
         * @date 2021/11/1
         */
        private AVLNode<K, V> rightRotate(AVLNode<K, V> cur) {
            AVLNode<K, V> left = cur.l;
            cur.l = left.r;
            left.r = cur;
            cur.h = Math.max((cur.l != null ? cur.l.h : 0), (cur.r != null ? cur.r.h : 0)) + 1;
            left.h = Math.max((left.l != null ? left.l.h : 0), (left.r != null ? left.r.h : 0)) + 1;
            return left;
        }

        /**
         * 左旋
         *
         * @param cur
         * @return
         */
        private AVLNode<K, V> leftRotate(AVLNode<K, V> cur) {
            AVLNode<K, V> right = cur.r;
            cur.r = right.l;
            right.l = cur;
            cur.h = Math.max((cur.l != null ? cur.l.h : 0), (cur.r != null ? cur.r.h : 0)) + 1;
            right.h = Math.max((right.l != null ? right.l.h : 0), (right.r != null ? right.r.h : 0)) + 1;
            return right;
        }

        private AVLNode<K, V> add(AVLNode<K, V> cur, K key, V value) {
            //如果等于null 说明当前节点没节点了 当前的key 加入到这个位置 返回当前节点
            if (cur == null) {
                return new AVLNode<>(key, value);
            } else if (cur.k.compareTo(key) > 0) {
                //递归接受l的值
                cur.l = add(cur.l, key, value);
            } else {
                //递归接受r的值
                cur.r = add(cur.r, key, value);
            }
            //计算高度。高端等于叶节点最大的高度+1
            cur.h = Math.max(cur.l == null ? 0 : cur.l.h, cur.r == null ? 0 : cur.r.h) + 1;
            //从最底层递归开始 检查平衡性
            AVLNode<K, V> head = maintain(cur);
            return head;
        }

        private AVLNode<K, V> delete(AVLNode<K, V> cur, K key) {
            return null;
        }

        /**
         * <p>
         * 平衡检测
         * </p>
         *
         * @param cur
         * @return com.lx.algorithm.Tree.AVLTree.AVLNode<K, V>
         * @Author zlx
         * @date 2021/11/1
         */
        private AVLNode<K, V> maintain(AVLNode<K, V> cur) {
            if (cur == null) {
                return null;
            }
            int leftHeight = cur.l != null ? cur.l.h : 0;
            int rightHeight = cur.r != null ? cur.r.h : 0;
            if (Math.abs(leftHeight - rightHeight) > 1) {
                if (leftHeight > rightHeight) {
                    int leftLeftHeight = cur.l != null && cur.l.l != null ? cur.l.l.h : 0;
                    int leftRightHeight = cur.l != null && cur.l.r != null ? cur.l.r.h : 0;
                    if (leftLeftHeight >= leftRightHeight) {
                        cur = rightRotate(cur);
                    } else {
                        cur.l = leftRotate(cur.l);
                        cur = rightRotate(cur);
                    }
                } else {
                    int rightLeftHeight = cur.r != null && cur.r.l != null ? cur.r.l.h : 0;
                    int rightRightHeight = cur.r != null && cur.r.r != null ? cur.r.r.h : 0;
                    if (rightRightHeight >= rightLeftHeight) {
                        cur = leftRotate(cur);
                    } else {
                        cur.r = rightRotate(cur.r);
                        cur = leftRotate(cur);
                    }
                }
            }
            return cur;
        }

        /**
         * 查找当前节点是否存在于树中
         * 如果存在返回当前节点的节点信息
         * 如果不存在插入节点的父节点信息
         *
         * @param key
         * @return
         */
        private AVLNode<K, V> findLastIndex(K key) {
            AVLNode<K, V> cur = root;
            AVLNode<K, V> pre = root;

            while (cur != null) {
                pre = cur;
                if (cur.k.compareTo(key) == 0) {
                    break;
                } else if (cur.k.compareTo(key) < 0) {
                    cur = cur.l;
                } else {
                    cur = cur.r;
                }
            }
            return pre;
        }


        /**
         * <p>
         * 方法用于返回大于或等于给定键元素(ele)的下键元素，否则返回null。
         * </p>
         *
         * @param key
         * @return com.lx.algorithm.Tree.AVLTree.AVLNode<K, V>
         * @Author zlx
         * @date 2021/11/1
         */
        private AVLNode<K, V> findLastNoSmallIndex(K key) {
            AVLNode<K, V> cur = root;
            AVLNode<K, V> pre = null;
            while (cur != null) {
                //如果相等 返回当前节点
                if (cur.k.compareTo(key) == 0) {
                    pre = cur;
                    break;
                    //如果当前节点大于key 说明在左边 把要返回的节点先跟循环节点都赋值 去左边找
                } else if (cur.k.compareTo(key) > 0) {
                    pre = cur;
                    cur = cur.l;
                } else {
                    //如果右边节点一直小于查询节点 说明不存在一个大于或者等于的节点
                    cur = cur.r;
                }
            }
            return pre;
        }

        /**
         * 方法用于返回的最大键小于或等于给定的键，或null，如果不存在这样的键
         *
         * @param key
         * @return
         */
        private AVLNode<K, V> findLastNoBigIndex(K key) {
            AVLNode<K, V> cur = root;
            AVLNode<K, V> pre = null;
            while (cur != null) {
                if (cur.k.compareTo(key) == 0) {
                    pre = cur;
                    break;
                } else if (cur.k.compareTo(key) < 0) {
                    pre = cur;
                    cur = cur.r;
                } else {
                    cur = cur.l;
                }
            }
            return pre;
        }

        public int size() {
            return size;
        }

        public boolean containsKey(K key) {
            if (key == null) {
                return false;
            }
            AVLNode<K, V> lastIndex = findLastIndex(key);
            if (lastIndex != null && lastIndex.k == key) {
                return true;
            }
            return false;
        }

        /**
         * 添加
         *
         * @param key
         * @param value
         */
        public void put(K key, V value) {
            if (key == null) {
                return;
            }
            //判断当前节点是否在树里存在
            AVLNode<K, V> lastIndex = findLastIndex(key);
            if (lastIndex != null && key.compareTo(lastIndex.k) == 0) {
                lastIndex.v = value;
            } else {
                root = add(root, key, value);
                size++;
            }
        }

        public void remove(K key) {
            if (key == null) {
                return;
            }
            if (containsKey(key)) {
                //一定要换头 因为删除一个节点以后可能存在头节点变换 所以要用返回的头代替原来的头
                root = delete(root, key);
                size--;
            }
        }

        /**
         * <p>
         * 如果存在key 返回value 不存在返回null
         * </p>
         *
         * @param key
         * @return V
         * @Author zlx
         * @date 2021/11/1
         */
        public V get(K key) {
            if (root == null) {
                return null;
            }
            AVLNode<K, V> lastIndex = findLastIndex(key);
            if (lastIndex != null && lastIndex.k == key) {
                return lastIndex.v;
            }
            return null;
        }

        /**
         * 方法用于返回当前这个映射第一个(最低)键。
         *
         * @return
         */
        public K firstKey() {
            if (root == null) {
                return null;
            }
            AVLNode<K, V> cur = root;
            while (cur.l != null) {
                cur = cur.l;
            }
            return cur.k;
        }

        /**
         * <p>
         * 方法用于返回当前在此映射的最后一个(最高)键。
         * </p>
         *
         * @param
         * @return K
         * @Author zlx
         * @date 2021/11/1
         */
        public K lastKey() {
            if (root == null) {
                return null;
            }
            AVLNode<K, V> cur = root;
            while (cur.r != null) {
                cur = cur.r;
            }
            return cur.k;
        }

        /**
         * 方法用于返回的最大键小于或等于给定的键，或null，如果不存在这样的键
         *
         * @param key
         * @return
         */
        public K floorKey(K key) {
            if (key == null) {
                return null;
            }
            AVLNode<K, V> lastNoBigIndex = findLastNoBigIndex(key);
            return lastNoBigIndex == null ? null : lastNoBigIndex.k;
        }

        /**
         * <p>
         * 方法用于返回大于或等于给定键元素(ele)的下键元素，否则返回null。
         * </p>
         *
         * @param key
         * @return K
         * @Author zlx
         * @date 2021/11/1
         */
        public K ceilingKey(K key) {
            if (key == null) {
                return null;
            }
            AVLNode<K, V> lastNoSmallNode = findLastNoSmallIndex(key);
            return lastNoSmallNode == null ? null : lastNoSmallNode.k;
        }
    }
}
